'''
FilePath: /src/utils/plot.py
Description: 画图

Author: xy
Date: 2025-06-11 15:42:50
Version: 
'''



# 标准库

# 第三方库
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches

# 自定义库

def plot_gantt_chart(machine_usage):
    
    '''
    description: 绘制甘特图
    param {*} machine_usage: 机器使用表
    '''

    
    # 创建图形
    plt.figure(figsize=(15, 8))
    
    # 按机器分组
    machine_data = {}
    for machine_id, usages in enumerate(machine_usage):
        for usage in usages:
            entry = {
                'Machine': machine_id,
                'Start': usage[0],
                'End': usage[1],
                'Heat': usage[2],  # 炉次号
                'Stage': usage[3]   # 阶段号
            }
            if machine_id not in machine_data:
                machine_data[machine_id] = []
            machine_data[machine_id].append(entry)
    
    # 定义颜色方案
    colors = plt.cm.tab20.colors
    
    # 为每个机器创建轨迹
    for y, machine in enumerate(sorted(machine_data.keys())):
        for entry in machine_data[machine]:
            plt.barh(y,
                    width=entry['End']-entry['Start'],
                    left=entry['Start'],
                    color=colors[entry['Stage'] % 20],
                    edgecolor='black')
            plt.text(entry['Start'] + 0.5, y,
                    f"{entry['Heat']}",
                    va='center', ha='left', fontsize=8)
    
    # 设置图表属性
    plt.yticks(range(len(machine_data)), [f'Machine {i}' for i in sorted(machine_data.keys())])
    plt.xlabel('Time')
    plt.title('Steel Cast Scheduling Gantt Chart')
    
    # 创建阶段图例
    stage_legend = [mpatches.Patch(color=colors[s%20], label=f'Stage {s}') 
                   for s in sorted({d['Stage'] for machine in machine_data.values() for d in machine})]
    plt.legend(handles=stage_legend, loc='upper right')
    plt.tight_layout()
    plt.show()